import Image, ImageFilter
from image_util import isBlack, RgbToLab, RgbToGrey, getPatch, getLabPatch, colorDiff
from ppm import read
from math import sqrt
from vec import dot, perp, norm, mag
import numpy as np

def getNeighbors(i, j):
	# assume (i,j) is not on image border
	n = []
	#n.append((i-1,j-1))
	n.append((i-1,j))
	#n.append((i-1,j+1))
	n.append((i,j-1))
	n.append((i,j+1))
	#n.append((i+1,j-1))
	n.append((i+1,j))
	#n.append((i+1,j+1))
	return n

def isBoundary(l, i, j):
	if not isBlack(l, i, j):
		return False
	for (x,y) in getNeighbors(i,j):
		if not isBlack(l,x,y):
			return True
	return False
	
def imageGradient(labMatrix, i, j):
	#print "Calculating image gradient at (%d,%d)" % (i,j)
	cx2 = labMatrix[i+1][j][0]
	#print "cx2:",cx2
	cx1 = labMatrix[i-1][j][0]
	#print "cx1:",cx1
	dx = (cx2 - cx1)/2.0
	#print "dx",dx
	
	cy2 = labMatrix[i][j+1][0]
	#print "cy2:",cy2
	cy1 = labMatrix[i][j-1][0]
	#print "cy1:",cy1 
	dy = (cy2 - cy1)/2.0
	#print "dy:",dy
	
	norm = sqrt(dx*dx+dy*dy)
	if norm == 0.0:
		return (0.0,0.0,0.0)
	else:
		return (dx/norm,dy/norm,0.0)
	
def normalGradient(im, i, j):
	# assume im is blurred
	cx2 = RgbToGrey(im.getpixel((i+1,j)))
	cx1 = RgbToGrey(im.getpixel((i-1,j)))
	dx	= (cx2 - cx1)/2.0

	cy2 = RgbToGrey(im.getpixel((i,j+1)))
	cy1 = RgbToGrey(im.getpixel((i,j-1))) 
	dy	= (cy2 - cy1)/2.0

	norm = sqrt(dx*dx+dy*dy)
	if norm == 0.0:
		return (0.0,0.0,0.0)
	else:
		return (dx/norm,dy/norm,0.0)

def initConfidenceMatrix(RGBMatrix):
	mat = []
	for i in range(len(RGBMatrix)):
		col = []
		for j in range(len(RGBMatrix[0])):
			if isBlack(RGBMatrix, i, j):
				col.append(0.0)
			else:
				col.append(1.0)
		mat.append(col)
	return mat
	
def confidenceTerm(mat, x, y, size=9):
	C = 0
	offset = (size - 1) / 2
	for i in range(x - offset, x + offset + 1):
		for j in range(y - offset, y + offset + 1):
			C += mat[i][j]
			
	return C / float(size * size)
			
def dataTerm(blurImage, gradientMatrix, x, y, patchSize):
	maxGradVal = 0.0
	maxGradient = (0.0,0.0,0.0) #might miss bugs with this
	
	offset = (patchSize-1)/2
	for i in range(x - offset, x + offset + 1):
		for j in range(y - offset, y + offset + 1):
			grad = mag(gradientMatrix[i][j])
			if grad > maxGradVal:
				maxGradVal = grad
				maxGradient = gradientMatrix[i][j]
				
	g = perp(maxGradient)
	n = norm(normalGradient(blurImage, x, y))
	return abs(dot(g, n))
		
def boundaryTest(filename="../asst2io/inpainting_major_clean.ppm"):
	v=5.0
	img = read(filename)
	blur_img = img.filter(ImageFilter.BLUR)
	newimg = Image.new("RGB", img.size)
	pix = newimg.load()
	width,height = img.size
	for y in range(1,height-1):
		for x in range(1,width-1):
			if isBoundary(img, x, y):
				pix[x,y] = (0,255,0)
			else:
				pix[x,y] = img.getpixel((x,y))
	for y in range(1,height-1):
		for x in range(1,width-1):
			if pix[x,y][1] == 255:
				grad = normalGradient(blur_img, x, y)
				pix[int(x+v*grad[0]), int(y+v*grad[1])] = (0,0,255)
	newimg.show()

def initPatches(mat, size = 9):
	offset = (size - 1) / 2
	size3 = size*size*3
	patches = []
	indexMap = []
	for i in range(offset, len(mat) - offset - 1):
		for j in range(offset, len(mat[0]) - offset - 1):
			patch = np.array(getLabPatch(mat,i,j,size))
			if len(patch.transpose().nonzero()[0]) == size3:
				patches.append(patch)
				indexMap.append((i,j))
			#	  print "Filled patch found at (%d, %d)" % (i, j)
			# else:
			#	  print "Unfilled patch found at (%d, %d)" % (i, j)
	yield patches
	yield indexMap
		
def findExemplar(mat, x, y, size = 9):
	offset = (size - 1) / 2
	p1 = getPatch(mat, x, y, size)
	currPatch = minimum = None
	
	for i in range(offset, len(mat) - offset - 1):
		for j in range(offset, len(mat[0]) - offset - 1):
			p2 = getPatch(mat, i, j, size)
			if filter(lambda x: (0,0,0) in x, p2) == []:
				diff = colorDiff(p1, p2)
				if minimum == None or diff < minimum:
					minimum = diff
					currPatch = p2
	
	return currPatch	